package com.hj.file.util;

import org.apache.lucene.analysis.Analyzer;
import org.apache.lucene.analysis.TokenStream;
import org.apache.lucene.document.*;
import org.apache.lucene.index.*;
import org.apache.lucene.queryparser.classic.QueryParser;
import org.apache.lucene.search.*;
import org.apache.lucene.search.highlight.*;
import org.apache.lucene.store.Directory;
import org.apache.lucene.store.FSDirectory;
import org.apache.poi.hwpf.extractor.WordExtractor;
import org.apache.poi.xwpf.extractor.XWPFWordExtractor;
import org.apache.poi.xwpf.usermodel.XWPFDocument;
import org.wltea.analyzer.lucene.IKAnalyzer;

import java.io.*;
import java.nio.file.Paths;
import java.util.*;

import static com.hj.file.util.POIUtils.mkDirs;

public class LuceneUtils {
    /**
     * 查询
     */
    public static List<Map<String, Object>> queryParser(String indexDir, String queryText, String fieldName, int pageSize) throws Exception {
        List<Map<String, Object>> list = new ArrayList<>();
        Directory directory = FSDirectory.open(Paths.get(indexDir));
        IndexReader ireader = DirectoryReader.open(directory);
        IndexSearcher isearcher = new IndexSearcher(ireader);
        Analyzer analyzer = new IKAnalyzer();
        QueryParser parser = new QueryParser(fieldName, analyzer);
//        parser.setDefaultOperator(QueryParser.Operator.AND);
        //parser.setAllowLeadingWildcard(true);
        Query query = parser.parse("\"" + queryText + "\"");
        // 返回前10条
        TopDocs tds = isearcher.search(query, pageSize);
        ScoreDoc[] hits = tds.scoreDocs;
        QueryScorer scorer = new QueryScorer(query);
        Fragmenter fragmenter = new SimpleSpanFragmenter(scorer);
        SimpleHTMLFormatter simpleHTMLFormatter = new SimpleHTMLFormatter("<b><font color='red'>", "</font></b>");
        Highlighter highlighter = new Highlighter(simpleHTMLFormatter, scorer);
        highlighter.setTextFragmenter(fragmenter);
        // Iterate through the results:
        for (int i = 0; i < hits.length; i++) {
            Map<String, Object> map = new HashMap<>();
            Document hitDoc = isearcher.doc(hits[i].doc);
            map.put("docId", hits[i].doc);
            for (int j = 0; j < hitDoc.getFields().size(); j++) {
                IndexableField indexableField = hitDoc.getFields().get(j);
                String name = indexableField.name();
                String value = indexableField.stringValue();
                if ("contents".equalsIgnoreCase(name)) {
                    TokenStream tokenStream = analyzer.tokenStream("contents", new StringReader(value));
                    /**
                     * getBestFragment方法用于输出摘要（即权重大的内容）
                     */
                    String formater = highlighter.getBestFragment(tokenStream, value);
                    System.out.println(formater);
                    map.put(indexableField.name(), formater);
                } else {
                    map.put(indexableField.name(), indexableField.stringValue());
                }
            }
            list.add(map);
        }
        ireader.close();
        directory.close();
        return list;
    }

    public static void doc2lucene(String wordPath, String fileName, String indexDir) throws IOException {
        mkDirs(wordPath);
        InputStream in = new FileInputStream(new File(wordPath + File.separator + fileName));
        WordExtractor wordExtractor = new WordExtractor(in);
        String contents = wordExtractor.getText();
        wordExtractor.close();
        storeLuceneIndex(createDocument(fileName, contents, wordPath), indexDir);
    }

    public static void docx2lucene(String wordPath, String fileName, String indexDir) throws IOException {
        mkDirs(wordPath);
        InputStream in = new FileInputStream(new File(wordPath + File.separator + fileName));
        XWPFWordExtractor xwpfWordExtractor = new XWPFWordExtractor(new XWPFDocument(in));
        // 创建Field对象，并放入doc对象中
        String contents = xwpfWordExtractor.getText();
        storeLuceneIndex(createDocument(fileName, contents, wordPath), indexDir);
    }

    public static List<Document> createDocument(String title, String contents, String path) {
        List<Document> documents = new ArrayList<>();
        Document doc = new Document();
        doc.add(new Field(LuceneConstant.TITLE, title, StringField.TYPE_STORED));
        doc.add(new Field(LuceneConstant.CONTENTS, contents, TextField.TYPE_STORED));
        doc.add(new Field(LuceneConstant.FILEPATH, path, StringField.TYPE_STORED));
        String createDate = DateTools.dateToString(new Date(), DateTools.Resolution.DAY);
        doc.add(new Field(LuceneConstant.CREATEDATE, createDate, StringField.TYPE_STORED));
        documents.add(doc);
        return documents;
    }

    /**
     * pdf 生成索引
     */
    public static void pdf2lucene(String pdfPath, String fileName, String indexDirectory) throws IOException {
        String contents = POIUtils.readPdfContent(pdfPath, fileName);
        storeLuceneIndex(createDocument(fileName, contents, pdfPath), indexDirectory);
    }

    public static void pdf2lucene(InputStream inputStream, String fileName, String pdfPath, String indexDirectory) throws IOException {
        String contents = POIUtils.readPdfContent(inputStream);
        storeLuceneIndex(createDocument(fileName, contents, pdfPath), indexDirectory);
    }

    /**
     * 存储索引
     */
    public static void storeLuceneIndex(List<Document> documents, String indexDirectory) throws IOException {
        IndexWriter indexWriter = getIndexWriter(indexDirectory);
        indexWriter.addDocuments(documents);
        indexWriter.commit();
        indexWriter.close();
    }

    public static IndexWriter getIndexWriter(String indexDir) throws IOException {
        Analyzer analyzer = new IKAnalyzer();
        Directory directory = FSDirectory.open(Paths.get(indexDir));
        IndexWriterConfig config = new IndexWriterConfig(analyzer);
        IndexWriter indexWriter = new IndexWriter(directory, config);
        return indexWriter;
    }
}
